import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rcParams
from sklearn.neighbors import NearestNeighbors
from sklearn.decomposition import PCA
import matplotlib
matplotlib.use('TkAgg')  # 切换到 TkAgg 后端

# 设置字体以支持中文字符
rcParams['font.sans-serif'] = ['SimHei']  # 使用SimHei字体
rcParams['axes.unicode_minus'] = False   # 修复负号显示问题


# 可视化结果
def visualize_results(dataset_name, X, labels, true_labels, rknn_dpc, k=None, r=None):
    # 可视化聚类结果（二维）
    if X.shape[1] == 2:
        plt.figure(figsize=(15, 5))

        # 真实标签
        plt.subplot(1, 3, 1)
        plt.scatter(X[:, 0], X[:, 1], c=true_labels, cmap='tab10', s=30, alpha=0.8)
        plt.title(f'{dataset_name} - 真实标签')
        plt.xlabel('X1')
        plt.ylabel('X2')

        # 最佳聚类结果
        plt.subplot(1, 3, 2)
        unique_labels = np.unique(labels)
        colors = plt.cm.tab10(np.linspace(0, 1, max(3, len(unique_labels))))

        for i, label in enumerate(unique_labels):
            if label != -1:
                plt.scatter(X[labels == label, 0], X[labels == label, 1],
                            c=[colors[label]], s=30, alpha=0.8, label=f'类簇{label}')

        # 标记类簇中心
        plt.scatter(X[rknn_dpc.centers, 0], X[rknn_dpc.centers, 1],
                    s=200, edgecolors='red', facecolors='none',
                    linewidths=2, marker='*', label='类簇中心')

        plt.title(f'{dataset_name} - 聚类结果\n(K={k}, R={r})')
        plt.xlabel('X1')
        plt.ylabel('X2')
        plt.legend()
        
        # TODO 计算局部密度和相对距离并绘制决策图
        density, delta = compute_density_delta(X, rknn_dpc.k)
        plot_decision_graph(density, delta, rknn_dpc.centers, f'RKNN-DPC 决策图 - {dataset_name} 数据集')

        plt.show()
    else:
        # 对于高维数据，使用PCA降维后可视化
        visualize_high_dim_results(dataset_name, X, labels, true_labels, rknn_dpc)

# 计算局部密度和相对距离
def compute_density_delta(X, k):
    n_samples = X.shape[0]
    X_normalized = (X - X.min(axis=0)) / (X.max(axis=0) - X.min(axis=0) + 1e-10)
    
    # 计算距离矩阵
    dist_matrix = np.zeros((n_samples, n_samples))
    for i in range(n_samples):
        for j in range(n_samples):
            dist_matrix[i, j] = np.sqrt(np.sum((X_normalized[i] - X_normalized[j])**2))
    
    # 使用K近邻计算局部密度
    nbrs = NearestNeighbors(n_neighbors=k+1, algorithm='auto').fit(X_normalized)
    _, indices = nbrs.kneighbors(X_normalized)
    knn_indices = indices[:, 1:k+1]  # 排除自身
    
    # 计算K近邻密度
    density = np.zeros(n_samples)
    for i in range(n_samples):
        knn_dist_sum = np.sum([dist_matrix[i, j] for j in knn_indices[i]])
        density[i] = np.exp(-knn_dist_sum)
    
    # 计算相对距离
    delta = np.zeros(n_samples)
    for i in range(n_samples):
        higher_density_points = np.where(density > density[i])[0]
        if len(higher_density_points) > 0:
            delta[i] = np.min(dist_matrix[i, higher_density_points])
        else:
            delta[i] = np.max(dist_matrix[i])
    
    return density, delta

# 可视化高维数据的聚类结果
def visualize_high_dim_results(dataset_name, X, labels, true_labels, rknn_dpc):
    # 使用PCA降维到2D
    pca = PCA(n_components=2)
    X_pca = pca.fit_transform(X)
    
    # 显示原始数据（带有真实标签，如果有）
    plt.figure(figsize=(10, 8))
    if true_labels is not None:
        plt.scatter(X_pca[:, 0], X_pca[:, 1], c=true_labels, cmap='tab10', s=30, alpha=0.8)
        plt.title(f'{dataset_name} 数据集 (PCA降维) - 真实标签')
    else:
        plt.scatter(X_pca[:, 0], X_pca[:, 1], c='gray', s=30, alpha=0.5)
        plt.title(f'{dataset_name} 数据集 (PCA降维)')
    plt.xlabel('主成分1')
    plt.ylabel('主成分2')
    plt.tight_layout()
    plt.savefig(f'results/{dataset_name}_original_pca.png')
    plt.show()
    
    # 显示聚类结果
    plt.figure(figsize=(10, 8))
    plt.scatter(X_pca[:, 0], X_pca[:, 1], c=labels, cmap='tab10', s=30, alpha=0.8)
    
    # 如果有中心点，也在PCA空间中标记它们
    if rknn_dpc.centers is not None:
        plt.scatter(X_pca[rknn_dpc.centers, 0], X_pca[rknn_dpc.centers, 1], 
                   s=200, edgecolors='red', facecolors='none', 
                   linewidths=2, marker='*', label='类簇中心')
    
    plt.title(f'RKNN-DPC 聚类结果 - {dataset_name} 数据集 (PCA降维)')
    plt.xlabel('主成分1')
    plt.ylabel('主成分2')
    plt.legend()
    plt.tight_layout()
    plt.savefig(f'results/{dataset_name}_clusters_pca.png')
    plt.show()
    
    # 计算并绘制决策图
    density, delta = compute_density_delta(X, rknn_dpc.k)
    plot_decision_graph(density, delta, rknn_dpc.centers, f'RKNN-DPC 决策图 - {dataset_name} 数据集')

# 绘制决策图函数
def plot_decision_graph(rho, delta, centers=None, title="RKNN-DPC 决策图"):
    """
    rho: 局部密度数组
    delta: 相对距离数组
    centers: 类簇中心点索引列表（可选）
    title: 图表标题
    """
    plt.subplot(1, 3, 3)
    # 创建颜色映射
    scatter = plt.scatter(rho, delta, c=np.arange(len(rho)), cmap='viridis', alpha=0.7)
    
    # 如果提供了中心点，突出显示它们
    if centers is not None:
        for i, center in enumerate(centers):
            plt.scatter(rho[center], delta[center], s=150, 
                       edgecolors='r', facecolors='none', linewidths=2,
                       marker='*')
    
    plt.xlabel('K近邻密度 ρ', fontsize=14)
    plt.ylabel('相对距离 δ', fontsize=14)
    plt.title(title, fontsize=16)
    plt.grid(True)
    plt.colorbar(scatter, label='点索引')

    plt.tight_layout()
    
    # 保存图像
    result_dir = 'results'
    if not os.path.exists(result_dir):
        os.makedirs(result_dir)
        
    filename = title.replace(' ', '_').replace('-', '_')
    plt.savefig(f'{result_dir}/{filename}.png')
    
    plt.show()
    
# 绘制聚类结果函数
def plot_clusters(X, labels, centers=None, title="RKNN-DPC 聚类结果"):
    """
    X: 数据集，可以是任意维度
    labels: 聚类标签
    centers: 类簇中心点索引列表（可选）
    title: 图表标题
    """
    # 确保在新窗口中绘制
    plt.figure(figsize=(10, 8))
    
    unique_labels = np.unique(labels)
    
    # 使用更丰富的颜色方案，支持更多类别
    if len(unique_labels) <= 10:
        colors = plt.cm.tab10(np.linspace(0, 1, len(unique_labels)))
    else:
        colors = plt.cm.viridis(np.linspace(0, 1, len(unique_labels)))
    
    # 对于二维数据，直接绘制散点图
    if X.shape[1] == 2:
        # 对于未分配的点（如果有）
        if -1 in unique_labels:
            mask = labels == -1
            plt.scatter(X[mask, 0], X[mask, 1], c='black', marker='x', label='未分配')
            unique_labels = unique_labels[unique_labels != -1]  # 移除-1
        
        # 绘制每个类簇
        for i, label in enumerate(unique_labels):
            mask = labels == label
            plt.scatter(X[mask, 0], X[mask, 1], c=[colors[i]], label=f'类簇 {label}')
            
        # 如果提供了中心点，突出显示它们
        if centers is not None:
            plt.scatter(X[centers, 0], X[centers, 1], s=200, 
                       edgecolors='red', facecolors='none', linewidths=2,
                       marker='*', label='类簇中心')
                        
        plt.xlabel('特征1', fontsize=14)
        plt.ylabel('特征2', fontsize=14)
    
    # 对于三维数据，绘制3D散点图
    elif X.shape[1] == 3:
        ax = plt.axes(projection='3d')
        
        # 对于未分配的点（如果有）
        if -1 in unique_labels:
            mask = labels == -1
            ax.scatter(X[mask, 0], X[mask, 1], X[mask, 2], 
                    c='black', marker='x', label='未分配')
            unique_labels = unique_labels[unique_labels != -1]  # 移除-1
        
        # 绘制每个类簇
        for i, label in enumerate(unique_labels):
            mask = labels == label
            ax.scatter(X[mask, 0], X[mask, 1], X[mask, 2], 
                    c=[colors[i]], label=f'类簇 {label}')
            
        # 如果提供了中心点，突出显示它们
        if centers is not None:
            ax.scatter(X[centers, 0], X[centers, 1], X[centers, 2], 
                    s=200, edgecolors='red', facecolors='none', 
                    linewidths=2, marker='*', label='类簇中心')
                        
        ax.set_xlabel('特征1', fontsize=14)
        ax.set_ylabel('特征2', fontsize=14)
        ax.set_zlabel('特征3', fontsize=14)
    
    else:
        # 高维数据应该已经通过visualize_high_dim_results处理
        print("警告：不支持直接绘制高维数据，请使用visualize_high_dim_results函数")
        return
    
    plt.title(title, fontsize=16)
    plt.legend()
    plt.tight_layout()
    
    # 保存图像
    result_dir = 'results'
    if not os.path.exists(result_dir):
        os.makedirs(result_dir)
        
    filename = title.replace(' ', '_').replace('-', '_')
    plt.savefig(f'{result_dir}/{filename}.png')
    
    plt.show()

